# an attempt at Lloyd's Algorithm
# for now, input points as an array of arrays
import random
import math

# test set:
points = [[1,6],[2,5],[8,9],[9,8],[7,10],[4,5],[9,4],[2,3],[7,7],[9,9]]


# takes an array of points and groups them into k clusters
def lloyd(pointsArray, k):
    dim = len(pointsArray[1])
    random.seed()
    # make a dictionary of cluster -> members
    # also a dictionary of cluster -> centroid
    cToMembers = {}
    cToCentroid = {}
    for i in range(k):
        cToMembers[i] = []
        cToCentroid[i] = [0,0,0]
        
    # initialize by randomly assigning each point to a cluster
    # make sure none contain 0 points
    random.shuffle(points)
    clumpSize = len(pointsArray)/k
    start = 0
    for cluster in range(k-1):
        cToMembers[cluster].extend(points[start:start+clumpSize])
        start += clumpSize
    cToMembers[k-1].extend(points[start:])
    print("initial clusters: " + str(cToMembers))
    done = False
    count = 0
    while not done:
        count += 1
        cToMembers, cToCenters, done = iterate(cToMembers, cToCentroid, pointsArray, dim, k)
    print('converged in ' + str(count) + ' iterations')
    return cToMembers, cToCenters

# this method takes in dicts of the cluster code to members and centroids
# as well as the original points array
# Python passes by reference, so no need to return anything
# 1. compute centroid of each cluster
# 2. reassign points to whichever centroid is closer
def iterate(cMembers, cCenters, pointsArray, dim, k):
    is_done = False

    # COMPUTE CENTROID OF EACH CLUSTER
    has_converged = True
    for cluster in cMembers.keys():
        oldCentroid = cCenters[cluster]
        numPoints = len(cMembers[cluster])
        centroid = []
        # initially populate centroid list with zeros
        for i in range(dim):
            centroid.insert(i,0)
        # use centroid list to keep totals, which we will average
        for member in cMembers[cluster]:
            for d in range(len(member)):
                centroid[d]+=member[d]
        # now compute averages
        for i in range(len(centroid)):
            centroid[i] /= float(numPoints)
            if not centroid[i] == oldCentroid[i]:
                has_converged = False
        # now shove this cluster into the other dictionary
        cCenters[cluster] = centroid
        
    
    # assign each point to the closest cluster (Euclidean distance)
    
    cMembers = {}
    for i in range(k):
        cMembers[i] = []
    
    for point in pointsArray:
        cToDistance = {}
        minDistance = float("inf")
        bestCluster = -1
        for c in cCenters.keys():
            ctr = cCenters[c]
            dsum = 0
            for d in range(dim):
                diff = ctr[d] - point[d]
                dsum += math.pow(diff,2)
            dist = math.sqrt(dsum)
            cToDistance[c] = dist
            if dist < minDistance:
                minDistance = dist
                bestCluster = c
        cMembers[bestCluster].append(point)
    return cMembers, cCenters, has_converged
    # actually return just the centroid of each cluster    
        
result, centroids = lloyd(points,2)
print("final clusters: " + str(result))
print("final centroids: " + str(centroids))
    
